import os
import gymnasium as gym
from stable_baselines3 import PPO

# 保存路径
SAVE_PATH = "ppo_lunarlander"
LOG_DIR = "ppo_log"

# 创建环境
env = gym.make("LunarLander-v3", render_mode="human")


# 如果已有模型文件，则加载继续训练
if os.path.exists(SAVE_PATH + ".zip"):
    print("🔄 检测到已有模型，加载并继续训练...")
    model = PPO.load(SAVE_PATH, env=env, tensorboard_log=LOG_DIR, device="cpu")
else:
    print("✨ 没有检测到模型，新建一个 PPO 模型...")
    model = PPO("MlpPolicy", env, verbose=1, tensorboard_log=LOG_DIR, device="cpu")

# 分批训练，每次 50k 步
TIMESTEPS = 50000
for i in range(50):   # 跑 10 次，共 50万步
    print(f"\n🚀 第 {i+1} 次训练开始（{TIMESTEPS} timesteps）...")
    model.learn(total_timesteps=TIMESTEPS, reset_num_timesteps=False)
    # 每次训练完保存一次
    model.save(SAVE_PATH)
    print(f"✅ 第 {i+1} 次训练结束，模型已保存到 {SAVE_PATH}.zip")

env.close()

